import chisel3._
import chisel3.util._

class WBU(xlen: Int) extends Module {
  val io = IO(new Bundle {
    val in = Flipped(Decoupled(new MessageMW(xlen)))
    val out = Decoupled(new MessageWI(xlen))
    val wb_sel  = Output(UInt(2.W))
    val pc      = Output(UInt(xlen.W))
    val wb_addr = Output(UInt(5.W))
    val rd_data = Output(UInt(xlen.W))
  })

  // pipepline registers (meu => wbu)
  val wbu_inst     = RegInit(0.U(xlen.W))
  val wbu_pc       = RegInit(0.U(xlen.W))
  val wbu_PC_sel   = RegInit(0.U(2.W))
  val wbu_wb_sel   = RegInit(0.U(2.W))
  val wbu_csr_cmd  = RegInit(0.U(3.W))
  val wbu_wb_en    = RegInit(0.B)
  val wbu_br_taken = RegInit(0.B) // NOTE: for simulation to identify next pc
  val wbu_Alu_out  = RegInit(0.U(xlen.W))
  val wbu_wb_addr  = RegInit(0.U(5.W))
  val wbu_rd_data  = RegInit(0.U(xlen.W))

  // valid for output
  val valid = RegInit(0.B)

  // ready for output
  val ready = RegInit(1.B)

  val csr = Module(new CSR(xlen))

  wbu_inst     := Mux(io.in.valid && io.in.ready, io.in.bits.inst,     wbu_inst)
  wbu_pc       := Mux(io.in.valid && io.in.ready, io.in.bits.pc,       wbu_pc)
  wbu_PC_sel   := Mux(io.in.valid && io.in.ready, io.in.bits.PC_sel,   wbu_PC_sel)
  wbu_wb_sel   := Mux(io.in.valid && io.in.ready, io.in.bits.wb_sel,   wbu_wb_sel)
  wbu_csr_cmd  := Mux(io.in.valid && io.in.ready, io.in.bits.csr_cmd,  wbu_csr_cmd)
  wbu_wb_en    := Mux(io.in.valid && io.in.ready, io.in.bits.wb_en,    wbu_wb_en)
  wbu_br_taken := Mux(io.in.valid && io.in.ready, io.in.bits.br_taken, wbu_br_taken) // NOTE: for simulation to identify next pc
  wbu_Alu_out  := Mux(io.in.valid && io.in.ready, io.in.bits.Alu_out,  wbu_Alu_out)
  wbu_wb_addr  := Mux(io.in.valid && io.in.ready, io.in.bits.wb_addr,  wbu_wb_addr)
  wbu_rd_data  := Mux(io.in.valid && io.in.ready, io.in.bits.rd_data,  wbu_rd_data)

  csr.io.inst     := wbu_inst
  csr.io.pc       := wbu_pc
  csr.io.csr_cmd  := wbu_csr_cmd
  csr.io.csr_data := wbu_Alu_out

  io.out.bits.PC_sel   := wbu_PC_sel
  io.wb_sel            := wbu_wb_sel
  io.out.bits.wb_en    := wbu_wb_en
  io.out.bits.csr_out  := csr.io.csr_out
  io.out.bits.Alu_out  := wbu_Alu_out
  io.wb_addr           := wbu_wb_addr
  io.rd_data           := wbu_rd_data
  io.pc                := wbu_pc

  valid := Mux(io.out.valid & io.out.ready , 0.B, Mux(io.in.valid & io.in.ready, 1.B, valid))
  io.out.valid := valid

  ready := Mux(io.in.valid & io.in.ready, 0.B, Mux(io.out.valid & io.out.ready, 1.B, ready))
  io.in.ready  := ready

  dontTouch(wbu_PC_sel)
  dontTouch(wbu_br_taken)
}
